'''
Total reward:  222000.0
Mean reward per episode:  2220.0
'''

import gymnasium as gym
import ale_py
from lib import common

gym.register_envs(ale_py)

# Initialize the environment
# env = RewardPenaltyWrapper(gym.make('ALE/BattleZone-v5', render_mode="human"))
# env = gym.make('ALE/BattleZone-v5')
env = common.ProcessFrame84(gym.make('ALE/BattleZone-v5', render_mode='human'))

# Reset the environment to get the initial state
state = env.reset()
total_reward = 0
# Run a loop to play the game
episode = 100
total_reward = 0
for _ in range(episode):
    env.reset()
    for _ in range(100000):
        # Take a random action
        action = env.action_space.sample()

        # Get the next state, reward, done flag, and info from the environment
        state, reward, done, trunc, info = env.step(action)
        if reward != 0:
            total_reward += reward
            print("reward: ", reward)
            print("info: ", info)

        # If done, reset the environment
        if done or trunc:
            break

print("Total reward: ", total_reward)
print("Mean reward per episode: ", total_reward / episode)

# Close the environment
env.close()